%% Use OptDMD method to extract coherent strcutre from the secondary flow
% Sijie Sun et. cl. PNAS, 2023
% Based on Travis Askham 2017, OptDMD method
%% Choose Dataset
close all;clc;clear
load('Velocity_Field.mat')
addpath('./src');
%% Setup
nt=size(U_tensor,1);%number of frames
nx=size(U_tensor,2);%x dimension
ny=size(U_tensor,3);%y dimension
dt=t(2)-t(1);
mm=1e-3;%1 mm = 1e-3 m
%% Preview Velocity Field
figure()    ;hold on
quiver(X/mm,Y/mm,squeeze(mean(U_tensor(:,:,:),1)),squeeze(mean(V_tensor(:,:,:),1)),2,'LineWidth',1.5)
axis equal
xlabel('x (mm)')
ylabel('y (mm)')
%%
stride = 3;
f=figure('Units',"inches",'Position',[1,1,1.8,1.8]);
quiver(X(1:stride:end, 1:stride:end), Y(1:stride:end, 1:stride:end), ...
       squeeze(U_tensor(1,1:stride:end, 1:stride:end)), squeeze(V_tensor(1,1:stride:end, 1:stride:end)),2,'LineWidth',1.5);
xlim([0,50]*mm);ylim([0,50]*mm);
axis equal
ax=gca;
ax.Position = [0 0 1 1];
axis off
saveas(gcf,['First Frame'],'epsc')

%%
% v = VideoWriter('Velocity Field.mp4','MPEG-4');
% v.FrameRate=20;v.Quality=100;
% open(v);
% fig=figure('units','pixels','position',[0 0 900 900]);
% for i=1:nt
%     U_plot=squeeze(U_tensor(i,:,:))/mm;
%     V_plot=squeeze(V_tensor(i,:,:))/mm;
%     U_plot(1,35)=30;
%     a=quiver(X/mm,Y/mm,U_plot,V_plot,2,'LineWidth',1.5);
%     a.Color=[0.2118    0.2706    0.3098];
%     set(fig,'color','w');
%     title(['t = ' num2str(t(i),'%5.2f') ' s'],'FontSize',20)
%     text(44.2,2.7, '30 mm/s', 'HorizontalAlignment', 'center','FontSize',12);
%     axis equal
%     axis off
%     drawnow
%     writeVideo(v,frame2im(getframe(fig)))
% end
% close(v);
% disp('COMPLETED');
%% Merge 2D Velocity field into one vector
u=zeros(nx*ny*2,nt);
for i=1:nt
    u(:,i)=UV2X(reshape(U_tensor(i,:,:),[],1),reshape(V_tensor(i,:,:),[],1));
end
%% SVD decomposition with truncation for reference
[U, S, V]=svd(u,'econ');
figure()
stem(diag(S).^2);% Importances of the first ten SVD modes
set(gca,'YScale','log')
xlim([0,10])
% truncation at r
r =21;
U = U(:,1:r);
S = S(1:r,1:r);
V = V(:,1:r);
uhat=U*S*V';
figure()
hold on
plot(t,mean(u,1),'r-','DisplayName','Ground Truth')
plot(t,mean(uhat,1),'b--','linewidth',2,'DisplayName','SVD with truncation')
xlim([0,30])
xlabel('Time-s')
ylabel('Speed-m/s')

legend('Location','best')
set(gca,'fontsize',16)

%% Optimized DMD, On truncated modes
r=21;
OPT.maxiter=2000;
[Phi_optDMD,eigs_optDMD,b_optDMD]=optdmd(u, t,r,2,OPT);
%% Compare decomposited flow with the measured flow

uxhat2=zeros(1,nt);
ux_true=zeros(1,nt);
for j=1:nt
    reconstructed=zeros(size(u(:,1)));
    for i=1:r
        reconstructed=reconstructed+b_optDMD(i)*exp(eigs_optDMD(i)*dt*j)*Phi_optDMD(:,i);
    end    
    uxhat2(j)=mean(real(reconstructed(1:2:end)));
    ux_true(j)=mean(u(1:2:end,j));
end


f=figure('Units',"inches",'Position',[1,1,1.8,1.4317]);

hold on
c=plot(t,uxhat2/mm,'DisplayName','Reconstructed Value','linewidth',2,'Color','#4daf4a');
a=plot(t,ux_true/mm,'DisplayName','Measured Value','linewidth',2,'Color','#377eb8');
a.Color=[a.Color,0.5];
c.Color=[c.Color,0.5];


hold off
legend('Location','best')
legend 'boxoff'
set(gca,'fontsize',10)
xlabel('Time-s')
ylabel('Frame Averaged Ux-mm/s')
xlim([0,30])
xticks([0:10:30]);
ylim([-2,5])
yticks([-2:2:4]);
box on
set(gca,'linewidth',1)
print(gcf,['DMD reconstruction Comparison'],'-dpng','-r500')
saveas(gcf,['DMD reconstruction'],'epsc')

%% Check L2 Difference
L2difference=0;
L2true=0;
for j=1:nt
    reconstructed=zeros(size(u(:,1)));
    for i=1:r
        reconstructed=reconstructed+b_optDMD(i)*exp(eigs_optDMD(i)*dt*j)*Phi_optDMD(:,i);
    end    
    L2difference=L2difference+norm(u(:,j)-reconstructed);
    L2true=L2true+norm(u(:,j));
end
disp(['relative L2 difference is ',num2str(L2difference/L2true),'L2 original is ',num2str(L2true)])

%% Plot  DMD spectrum
f2=figure();
hold on

theta = (0:1:100)*2*pi/100;
plot(cos(theta),sin(theta),'k--') % plot unit circle
scatter(real(exp(eigs_optDMD*dt)),imag(exp(eigs_optDMD*dt)),'ok')
axis([-1.1 1.1 -1.1 1.1]);
xlabel('Re of Eigen Value')
ylabel('Im of Eigen Value')
title('DMD Eigen Value Distribution')
axis equal
print(f2,['Spectrum of Modes'],'-dpng','-r500')

%% Plot DMD Mode Importance Distribution
f3=figure();
[Contribution,Order]=sort(abs(b_optDMD),'descend');
bar(Contribution(1:11)/sum(Contribution),'linewidth',0.0001)
ylim([5e-3,1])
set(gca,'yscale','log')
box on; grid on
xlabel('Mode Number')
ylabel('Contribution')
title('Mode Contribution Distribution')
%% Plot Flow structure Importance Distribution
Y1=[0;Contribution];Y2=zeros(11,0);
for i=1:11
    Y2(i)=sum(Y1(i*2-1:i*2));
end
f4=figure();
bar(Y2/sum(Y2),'linewidth',0.0001)
ylim([1e-3,1])
set(gca,'yscale','log')
box on; grid off
xlabel('Structure Number')
ylabel('Contribution')
title('Structure contribution Distribution')
set(gca,'fontsize',14)
set(gca,'linewidth',2)
print(f4,['Structure Contribution Distribution'],'-dpng','-r500')

%% Plot the eigen modes
Scale=3;
for j=1:2:5
    i=Order(j);
    [f,ax1,ax2]=plotVectorfield(Phi_optDMD(:,i),nx,ny,exp(eigs_optDMD(i)*dt),dt,Scale,X,Y,mm);
    print(f,['Velocity field of Sorted Mode ' num2str(j)],'-dpng','-r300')
    saveas(f,['Velocity field of Sorted Mode ' num2str(j)],'epsc')
end

%% Auxiliary Functions

function X=UV2X(U,V)
    X=zeros(length(U)*2,1);
    X(1:2:end-1)=U;
    X(2:2:end)=V;
end

function [f,ax1,ax2]=plotVectorfield(u,nx,ny,eig_val,dt,Scale,X,Y,mm)
    Uxr=reshape(real(u(1:2:end-1)),nx,ny);
    Uyr=reshape(real(u(2:2:end)),nx,ny);
    Uxi=reshape(imag(u(1:2:end-1)),nx,ny);
    Uyi=reshape(imag(u(2:2:end)),nx,ny);    
    Lr=norm(real(u(:)));
    Li=norm(imag(u(:)));
    f=figure('Position',[100,100,1600,700]);
    ax1=subplot(1,2,1);
        a=quiver(X/mm,Y/mm,Uxr/Lr,Uyr/Lr,Scale,'LineWidth',1.5);
        a.Color=[99, 109, 111]/255;
        axis equal
        title('Real Part')
        xlabel('x (mm)')
        ylabel('y (mm)')
        set(gca,'linewidth',2)
        set(gca,'fontsize',16)    
    ax2=subplot(1,2,2);
        a=quiver(X/mm,Y/mm,Uxi/Li,Uyi/Li,Scale,'LineWidth',1.5);
        a.Color=[99, 109, 111]/255;
        axis equal
        xlabel('x (mm)')
        ylabel('y (mm)')
        set(gca,'linewidth',2)
        set(gca,'fontsize',16)
        title('Imaginary Part')
    dtheta=angle(eig_val);
    omega=dtheta/dt;
    tau=2*pi/omega;
    sgtitle(['% \lambda = ',num2str(eig_val) ,' \tau =' ,num2str(abs(tau)),' s'],'fontsize',20) 
end


